import idc
import idaapi
import idautils
import ida_segment
import json

def set_image_base(new_base):
    # 获取当前图像基地址
    current_base = idaapi.get_imagebase()
    print(f"当前基地址: {hex(current_base)}")
    
    # 计算偏移量
    delta = new_base - current_base
    print(f"新基地址: {hex(new_base)}, 偏移量: {hex(delta)}")
    
    if delta == 0:
        print("基地址相同，无需重新基址。")
        return
    
    # 重新基址程序
    try:
        ida_segment.rebase_program(delta, ida_segment.MSF_FIXONCE)
        print(f"图像基地址已成功设置为 {hex(new_base)}")
    except Exception as e:
        print(f"设置图像基地址时出错: {e}")
        
# 获取汇编信息并格式化
def get_disasm_info(addr):
    try:
        disasm = idc.GetDisasm(addr).strip()
        if disasm:
            # 提取反汇编指令中的助记符和操作数
            disasm_parts = disasm.split(' ', 1)
            mnemonic = disasm_parts[0]
            op_str = disasm_parts[1] if len(disasm_parts) > 1 else ""
            
            return ({
                "address": hex(addr),
                "mnemonic": mnemonic,
                "op_str": op_str.strip()
            })
    except Exception as e:
        print(f"无法获取汇编信息: {e} (地址: {hex(addr)})")
    return None

def extract_callgraph(out_folder):
    nodes = []
    edges = []
    func_addr_to_id = {}
    node_id = 0

    unique_nodes = set()  # 用于去重节点
    unique_edges = set()  # 用于去重边 (src_addr, dst_addr)

    # 获取所有函数的节点信息
    for func_addr in idautils.Functions():
        func = idaapi.get_func(func_addr)
        # print(func_addr)
        if func:
            func_name = idc.get_func_name(func_addr) if func else f"sub_{func_addr:x}"
            # print(func_name)
            disasm_info = []

            # 获取函数的第一个基本块
            flowchart = idaapi.FlowChart(func)
            first_block = next(iter(flowchart), None)
            for addr in range(first_block.start_ea, first_block.end_ea):
                item_size = idc.get_item_size(addr)
                disasm_info.append(get_disasm_info(addr))
                addr += item_size

            node_info = {
                'id': func_addr,
                'addr': hex(func_addr),
                'function_name': func_name,
                'function_addr': hex(func_addr),
                'block_addr': hex(func_addr),  # 假设CG函数地址即为基本块地址
                'disasm': disasm_info
            }

            nodes.append(node_info)
            func_addr_to_id[func_addr] = node_id
            node_id += 1
            unique_nodes.add(func_addr)

    # 获取函数间调用的边信息
    # print(func_addr_to_id.keys())
    # for xref in idautils.XrefsTo(134516460,0):  # True 表示只获取控制流引用
    #     print(f"Reference found at: {hex(xref.frm)}")
    for dst_addr in func_addr_to_id.keys():
        for xref in idautils.XrefsTo(dst_addr, True):  # 获取从函数的引用
            # print(xref)
            if idaapi.get_func(xref.frm):
                src_addr = idaapi.get_func(xref.frm).start_ea
                # print(dst_addr)
                # print(func_addr_to_id.keys())
                if src_addr in func_addr_to_id.keys():
                    edge_key = (src_addr, dst_addr)  # 使用 (src_addr, dst_addr) 作为边的唯一标识
                    # print(edge_key)
                    if edge_key not in unique_edges:
                        edge_info = {
                            'src_id': src_addr,
                            'dst_id': dst_addr,
                            'type': "Call_Function"  # 边默认属性为 Call_Function
                        }
                        edges.append(edge_info)
                        unique_edges.add(edge_key)  # 将边添加到去重集合中

    # 对节点按照地址进行排序
    sorted_nodes = sorted(nodes, key=lambda x: x['id'])

    # 输出到文件
    output_path = f"{out_folder}/cg_graph.json"
    with open(output_path, 'w') as f:
        json.dump({'nodes': sorted_nodes, 'edges': edges}, f, indent=2)

    print(f"CG 已成功导出到 {output_path}")
    return sorted_nodes, edges

if __name__ == "__main__":
    idaapi.autoWait()
    out_folder = idc.ARGV[1]  # 从命令行参数获取输出目录
    current_base = idaapi.get_imagebase()
    print(f"当前基地址: {hex(current_base)}")

    # 定义新的基地址
    desired_base = 0x400000  # 设置您希望的基地址

    # 检测是否需要设置新的基地址
    if current_base == 0x0:
        print("检测到当前基地址为 0x0，准备设置新的基地址。")
        set_image_base(desired_base)
    else:
        print(f"当前基地址已设置为 {hex(current_base)}，无需重新基址。")
    extract_callgraph(out_folder)
    idc.qexit(0)